import json
import datetime
import json
import time

import tushare as ts
from dashscope import Generation
from tavily import TavilyClient

# ⚠️ 修复：彻底删除有问题的导入 ToolCall
# from dashscope.api_entities.dashscope_response import ToolCall
# 修复：导入新的 DDGS 库 (如果 Tavily 搜索失效，可以作为备用)
# from ddgs import DDGS

# ==========================================
# 1. 配置 Qwen, Tushare, Tavily Keys
# ==========================================

# 1. Qwen/DashScope Key
DASHSCOPE_API_KEY = "sk-f3a4b2876500475fb94b81fe9f37bd44"

# 2. ❗ 请替换为你自己的 Tushare Token
TS_TOKEN = "8c82bedea9058b9213d0402d5f1f10e2ad2d2f2d1177209e11e05481"
# 3. ❗ 请替换为你自己的 Tavily API Key
TAVILY_API_KEY = "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"

# 初始化 Tushare 客户端
ts.set_token(TS_TOKEN)
TS_PRO = ts.pro_api()

# 初始化 Tavily 客户端
TAVILY_CLIENT = TavilyClient(api_key=TAVILY_API_KEY)


# --- 工具函数定义：Tushare ---

def get_daily_k_data(ts_code: str):
    """
    输入股票的Tushare代码 (例如 '000001.SZ', '600519.SH')，返回最近一个交易日的收盘价、成交量和涨跌额。
    """
    print(f"\n[工具执行] 正在使用 Tushare 查询 {ts_code} 的最新K线数据...")
    time.sleep(1)

    today = datetime.date.today().strftime('%Y%m%d')

    try:
        df = TS_PRO.query(
            'daily',
            ts_code=ts_code,
            end_date=today,
            limit=10
        )

        if df.empty:
            return f"错误：Tushare找不到股票代码 {ts_code} 的数据，请检查代码是否正确或今日是否休市。"

        latest = df.iloc[0]

        return json.dumps({
            "ts_code": ts_code,
            "trade_date": latest['trade_date'],
            "close_price": float(latest['close']),
            "change_amount": float(latest['change']),
            "volume": float(latest['vol']) / 10000.0,
            "turnover_rate": float(latest['turnover_rate']) if 'turnover_rate' in df.columns else None,
            "analysis": f"最新收盘价: {latest['close']}，涨跌额: {latest['change']}，成交量: {latest['vol']}手。"
        })
    except Exception as e:
        return f"Tushare获取数据时发生错误: {e}"

# --- 工具函数定义：Tavily ---

def search_market_news(query: str):
    """
    搜索关于特定股票或市场的最新财经资讯。
    输入应为具体的搜索关键词，如 'NVIDIA 最新财报分析' 或 '上证指数行情'。
    返回前三条相关资讯的标题和摘要。
    """
    print(f"\n[工具执行] 正在使用 Tavily 搜索资讯: {query} ...")
    time.sleep(1)

    try:
        response = TAVILY_CLIENT.search(
            query=query,
            search_depth="basic",
            max_results=3
        )

        news_list = []
        for result in response['results']:
            news_list.append({
                "title": result.get("title"),
                "content": result.get("content"),
                "url": result.get("url")
            })

        return json.dumps(news_list, ensure_ascii=False)
    except Exception as e:
        return f"Tavily搜索资讯时发生错误: {e}"


# 工具映射表 (更新)
tool_map = {
    "get_daily_k_data": get_daily_k_data,
    "search_market_news": search_market_news,
}

# --- 工具描述列表 (Qwen 格式) ---
tools_schema = [
    {
        "type": "function",
        "function": {
            "name": "get_daily_k_data",
            "description": "输入股票的Tushare代码 (例如 '000001.SZ', '600519.SH')，获取最近一个交易日的收盘价、成交量和涨跌额等核心日K数据。",
            "parameters": {
                "type": "object",
                "properties": {
                    "ts_code": {"type": "string", "description": "股票的Tushare代码，如 '000001.SZ'"},
                },
                "required": ["ts_code"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "search_market_news",
            "description": "搜索关于特定股票或市场的最新财经资讯，返回3条摘要。",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "具体的搜索关键词，如 'NVIDIA 最新财报'"},
                },
                "required": ["query"],
            },
        },
    }
]

# ==========================================
# 3. Agent 核心逻辑 (Qwen 循环 - 稳定版)
# ==========================================

def run_stock_agent(prompt: str):
    messages = [
        {"role": "system", "content": "你是一位专业的金融分析师。请利用你拥有的Tushare数据和Tavily资讯工具，先查询股票数据和市场资讯，然后基于工具返回的所有信息，对股票进行全面的技术分析和基本面分析，给出专业的投资建议"},
        {"role": "user", "content": prompt}
    ]

    for i in range(5):
        print(f"\n--- 第 {i+1} 轮推理 ---")

        response = Generation.call(
            model='qwen-turbo',
            api_key=DASHSCOPE_API_KEY,
            messages=messages,
            tools=tools_schema,
        )

        if response.status_code != 200:
            return f"Qwen API 调用失败，状态码: {response.status_code}，原因: {response.message}"

        # VITAL FIX: 强制将 DashScope 对象转换为 Python 字典，消除所有版本兼容性问题
        response_message_json = json.dumps(response.output.choices[0].message)
        response_message_dict = json.loads(response_message_json)

        messages.append(response_message_dict)

        tool_calls = response_message_dict.get('tool_calls')

        if tool_calls:
            print(f"模型思考: 决定调用 {len(tool_calls)} 个工具...")

            tool_outputs = []

            for tool_call in tool_calls:
                function_name = tool_call['function']['name']
                function_args = tool_call['function']['arguments']

                try:
                    function_args_dict = json.loads(function_args)
                except json.JSONDecodeError:
                    print("Warning: Qwen返回的参数不是有效的JSON。")
                    function_args_dict = {}

                function_to_call = tool_map.get(function_name)
                if function_to_call:
                    tool_result = function_to_call(**function_args_dict)
                else:
                    tool_result = f"错误：未知的工具名称: {function_name}"

                tool_outputs.append({
                    "tool_call_id": tool_call['id'],
                    "output": tool_result,
                })

                print(f"[工具结果] {function_name} 返回长度: {len(tool_result)} 字符")

            messages.append(
                {
                    "role": "tool",
                    "tool_calls": tool_calls,
                    "content": json.dumps(tool_outputs, ensure_ascii=False)
                }
            )
        else:
            return response_message_dict.get('content')

    return "Agent 未能在 5 轮内完成分析，请检查模型或请求。"


# ==========================================
# 4. 运行 Agent
# ==========================================

if __name__ == "__main__":
    print("--- Qwen Agent (Tavily + Tushare) 已启动 ---")

    # ⚠️ 启动前检查 Key
    if "YOUR_TUSHARE_TOKEN_HERE" in TS_TOKEN or "YOUR_TAVILY_API_KEY_HERE" in TAVILY_API_KEY:
        print("\n!!! 错误: 请先将代码中的 Tushare Token 和 Tavily Key 替换为真实值 !!!\n")
        exit()

    stock_symbol = input("请输入你想分析的股票代码 (例如 000001.SZ, 600519.SH): ")

    if not stock_symbol:
        stock_symbol = "600519.SH" # 默认使用贵州茅台作为测试

    user_query = f"请分析 {stock_symbol} 的最新股价数据和市场资讯，并告诉我现在的投资建议。"

    print(f"\n开始分析 {stock_symbol}...")
    final_report = run_stock_agent(user_query)

    print("\n\n================ 最终分析报告 ================\n")
    print(final_report)
    print("\n==============================================\n")